# Copyright (c) Opendatalab. All rights reserved.
"""
本脚本提供两种后端（pipeline / vlm）对 PDF/图片进行解析，自动生成：
  - 结构化 Markdown
  - 中间 JSON
  - 模型原始输出
  - 版面/span 可视化 PDF
  - 原始 PDF 备份
  - 内容列表 JSON

使用方式：
  1. 把待解析文件放到与脚本同级的 `pdfs/` 目录
  2. 运行本脚本即可在 `output/` 目录获得结果
  3. 如需切换后端，修改最下方 parse_doc 的 backend 参数即可
"""

import copy
import json
import os
from pathlib import Path

from loguru import logger   # 日志库，打印进度与异常

# mineru 内部工具
from mineru.cli.common import (
    convert_pdf_bytes_to_bytes_by_pypdfium2,  # 将 PDF 按页码切片
    prepare_env,                              # 为每个 PDF 准备输出目录
    read_fn                                   # 统一读取文件（支持本地/PDF/图片）
)
from mineru.data.data_reader_writer import FileBasedDataWriter   # 写本地文件
from mineru.utils.draw_bbox import draw_layout_bbox, draw_span_bbox  # 可视化
from mineru.utils.enum_class import MakeMode   # 枚举：生成 Markdown 的几种模式
from mineru.backend.vlm.vlm_analyze import doc_analyze as vlm_doc_analyze
from mineru.backend.pipeline.pipeline_analyze import doc_analyze as pipeline_doc_analyze
from mineru.backend.pipeline.pipeline_middle_json_mkcontent import union_make as pipeline_union_make
from mineru.backend.pipeline.model_json_to_middle_json import result_to_middle_json as pipeline_result_to_middle_json
from mineru.backend.vlm.vlm_middle_json_mkcontent import union_make as vlm_union_make
from mineru.utils.models_download_utils import auto_download_and_get_model_root_path  # 首次运行时自动下载模型


def do_parse(
    output_dir,                     # 结果统一输出根目录
    pdf_file_names: list[str],      # 每个 PDF 的文件名（无后缀）
    pdf_bytes_list: list[bytes],    # 每个 PDF 的原始字节
    p_lang_list: list[str],         # 每个 PDF 的语言提示（提高 OCR 精度）
    backend="pipeline",             # 选择后端：pipeline / vlm-transformers / vlm-sglang-*
    parse_method="auto",            # 仅 pipeline 生效：auto/txt/ocr
    formula_enable=True,            # 是否识别公式
    table_enable=True,              # 是否识别表格
    server_url=None,                # vlm-sglang-client 后端需要填 server 地址
    f_draw_layout_bbox=True,        # 是否生成版面可视化 PDF
    f_draw_span_bbox=True,          # 是否生成 span 可视化 PDF
    f_dump_md=True,                 # 是否写出 .md
    f_dump_middle_json=True,        # 是否写出 _middle.json
    f_dump_model_output=True,       # 是否写出模型原始输出
    f_dump_orig_pdf=True,           # 是否备份原始 PDF
    f_dump_content_list=True,       # 是否写出 _content_list.json
    f_make_md_mode=MakeMode.MM_MD,  # Markdown 生成策略
    start_page_id=0,                # 从第几页开始解析
    end_page_id=None,               # 解析到第几页（None=全部）
):
    """
    真正执行解析的主函数，根据 backend 不同走两套流程。
    """

    # -----------------------------
    # 1. pipeline 后端
    # -----------------------------
    if backend == "pipeline":
        # 1.1 先把每份 PDF 按 start/end_page_id 裁剪
        for idx, pdf_bytes in enumerate(pdf_bytes_list):
            new_pdf_bytes = convert_pdf_bytes_to_bytes_by_pypdfium2(
                pdf_bytes, start_page_id, end_page_id
            )
            pdf_bytes_list[idx] = new_pdf_bytes

        # 1.2 统一送进 pipeline 解析，返回：
        #     infer_results: 每份 PDF 的模型输出
        #     all_image_lists: 每份 PDF 每页渲染出的图片路径
        #     all_pdf_docs: 每份 PDF 的 pypdfium2 对象
        #     lang_list/ocr_enabled_list: 每份 PDF 的语言及是否启用 OCR
        infer_results, all_image_lists, all_pdf_docs, lang_list, ocr_enabled_list = pipeline_doc_analyze(
            pdf_bytes_list,
            p_lang_list,
            parse_method=parse_method,
            formula_enable=formula_enable,
            table_enable=table_enable,
        )

        # 1.3 逐份 PDF 处理
        for idx, model_list in enumerate(infer_results):
            model_json = copy.deepcopy(model_list)                # 备份模型原始输出
            pdf_file_name = pdf_file_names[idx]                   # 当前 PDF 名字
            local_image_dir, local_md_dir = prepare_env(
                output_dir, pdf_file_name, parse_method
            )                                                     # 创建输出目录
            image_writer = FileBasedDataWriter(local_image_dir)   # 写图片
            md_writer = FileBasedDataWriter(local_md_dir)         # 写文本

            images_list = all_image_lists[idx]                    # 当前 PDF 的图片列表
            pdf_doc = all_pdf_docs[idx]                           # 当前 PDF 的对象
            _lang = lang_list[idx]
            _ocr_enable = ocr_enabled_list[idx]

            # 将模型输出转成统一中间格式
            middle_json = pipeline_result_to_middle_json(
                model_list, images_list, pdf_doc,
                image_writer, _lang, _ocr_enable, formula_enable
            )
            pdf_info = middle_json["pdf_info"]

            pdf_bytes = pdf_bytes_list[idx]

            # 1.3.1 可视化
            if f_draw_layout_bbox:
                draw_layout_bbox(
                    pdf_info, pdf_bytes,
                    local_md_dir, f"{pdf_file_name}_layout.pdf"
                )
            if f_draw_span_bbox:
                draw_span_bbox(
                    pdf_info, pdf_bytes,
                    local_md_dir, f"{pdf_file_name}_span.pdf"
                )

            # 1.3.2 备份原 PDF
            if f_dump_orig_pdf:
                md_writer.write(f"{pdf_file_name}_origin.pdf", pdf_bytes)

            # 1.3.3 生成 Markdown
            if f_dump_md:
                image_dir = str(os.path.basename(local_image_dir))
                md_content_str = pipeline_union_make(
                    pdf_info, f_make_md_mode, image_dir
                )
                md_writer.write_string(f"{pdf_file_name}.md", md_content_str)

            # 1.3.4 生成内容列表 JSON
            if f_dump_content_list:
                image_dir = str(os.path.basename(local_image_dir))
                content_list = pipeline_union_make(
                    pdf_info, MakeMode.CONTENT_LIST, image_dir
                )
                md_writer.write_string(
                    f"{pdf_file_name}_content_list.json",
                    json.dumps(content_list, ensure_ascii=False, indent=4),
                )

            # 1.3.5 导出中间 JSON
            if f_dump_middle_json:
                md_writer.write_string(
                    f"{pdf_file_name}_middle.json",
                    json.dumps(middle_json, ensure_ascii=False, indent=4),
                )

            # 1.3.6 导出模型原始输出
            if f_dump_model_output:
                md_writer.write_string(
                    f"{pdf_file_name}_model.json",
                    json.dumps(model_json, ensure_ascii=False, indent=4),
                )

            logger.info(f"local output dir is {local_md_dir}")

    # -----------------------------
    # 2. vlm-* 后端
    # -----------------------------
    else:
        # 2.1 去掉前缀 "vlm-" 得到真正的 backend 名称
        if backend.startswith("vlm-"):
            backend = backend[4:]

        # vlm 后端目前不支持 span 可视化
        f_draw_span_bbox = False
        parse_method = "vlm"

        # 2.2 逐份 PDF 处理
        for idx, pdf_bytes in enumerate(pdf_bytes_list):
            pdf_file_name = pdf_file_names[idx]
            # 裁剪页码
            pdf_bytes = convert_pdf_bytes_to_bytes_by_pypdfium2(
                pdf_bytes, start_page_id, end_page_id
            )
            # 准备输出目录
            local_image_dir, local_md_dir = prepare_env(
                output_dir, pdf_file_name, parse_method
            )
            image_writer = FileBasedDataWriter(local_image_dir)
            md_writer = FileBasedDataWriter(local_md_dir)

            # 调用 vlm 统一接口
            middle_json, infer_result = vlm_doc_analyze(
                pdf_bytes,
                image_writer=image_writer,
                backend=backend,
                server_url=server_url,
            )
            pdf_info = middle_json["pdf_info"]

            # 2.2.1 可视化
            if f_draw_layout_bbox:
                draw_layout_bbox(
                    pdf_info, pdf_bytes,
                    local_md_dir, f"{pdf_file_name}_layout.pdf"
                )
            if f_draw_span_bbox:
                draw_span_bbox(
                    pdf_info, pdf_bytes,
                    local_md_dir, f"{pdf_file_name}_span.pdf"
                )

            # 2.2.2 备份原 PDF
            if f_dump_orig_pdf:
                md_writer.write(f"{pdf_file_name}_origin.pdf", pdf_bytes)

            # 2.2.3 生成 Markdown
            if f_dump_md:
                image_dir = str(os.path.basename(local_image_dir))
                md_content_str = vlm_union_make(
                    pdf_info, f_make_md_mode, image_dir
                )
                md_writer.write_string(f"{pdf_file_name}.md", md_content_str)

            # 2.2.4 生成内容列表 JSON
            if f_dump_content_list:
                image_dir = str(os.path.basename(local_image_dir))
                content_list = vlm_union_make(
                    pdf_info, MakeMode.CONTENT_LIST, image_dir
                )
                md_writer.write_string(
                    f"{pdf_file_name}_content_list.json",
                    json.dumps(content_list, ensure_ascii=False, indent=4),
                )

            # 2.2.5 导出中间 JSON
            if f_dump_middle_json:
                md_writer.write_string(
                    f"{pdf_file_name}_middle.json",
                    json.dumps(middle_json, ensure_ascii=False, indent=4),
                )

            # 2.2.6 导出模型原始输出（文本拼接）
            if f_dump_model_output:
                model_output = ("\n" + "-" * 50 + "\n").join(infer_result)
                md_writer.write_string(
                    f"{pdf_file_name}_model_output.txt",
                    model_output,
                )

            logger.info(f"local output dir is {local_md_dir}")


def parse_doc(
    path_list: list[Path],
    output_dir,
    lang="ch",
    backend="pipeline",
    method="auto",
    server_url=None,
    start_page_id=0,
    end_page_id=None,
):
    """
    对外暴露的便捷函数，把「文件路径」转成「字节」后调用 do_parse。

    参数说明：
        path_list: 待解析文件路径列表（PDF/图片均可）
        output_dir: 结果输出根目录
        lang: 语言提示，pipeline 后端有效
        backend: pipeline | vlm-transformers | vlm-sglang-engine | vlm-sglang-client
        method: 仅 pipeline 有效：auto/txt/ocr
        server_url: vlm-sglang-client 时必须填写，如 http://127.0.0.1:30000
        start_page_id/end_page_id: 解析起始/终止页码
    """
    try:
        file_name_list = []
        pdf_bytes_list = []
        lang_list = []
        for path in path_list:
            file_name = str(Path(path).stem)   # 去掉扩展名
            pdf_bytes = read_fn(path)          # 读取文件为 bytes
            file_name_list.append(file_name)
            pdf_bytes_list.append(pdf_bytes)
            lang_list.append(lang)

        # 调真正的解析入口
        do_parse(
            output_dir=output_dir,
            pdf_file_names=file_name_list,
            pdf_bytes_list=pdf_bytes_list,
            p_lang_list=lang_list,
            backend=backend,
            parse_method=method,
            server_url=server_url,
            start_page_id=start_page_id,
            end_page_id=end_page_id,
        )
    except Exception as e:
        logger.exception(e)   # 发生异常时打印完整 traceback


# -----------------------------------------------------------------------------
# 脚本直接运行时的默认示例
# -----------------------------------------------------------------------------
if __name__ == '__main__':
    # 当前脚本所在目录
    __dir__ = os.path.dirname(os.path.abspath(__file__))

    # 1. 准备输入/输出目录
    pdf_files_dir = os.path.join(__dir__, "pdfs")
    output_dir = os.path.join(__dir__, "output")
    pdf_suffixes = [".pdf"]
    image_suffixes = [".png", ".jpeg", ".jpg"]

    # 2. 扫描待解析文件
    doc_path_list = []
    for doc_path in Path(pdf_files_dir).glob('*'):
        if doc_path.suffix.lower() in pdf_suffixes + image_suffixes:
            doc_path_list.append(doc_path)

    # 3. 如果因网络问题无法下载模型，可启用下面这行，使用 modelscope 镜像
    # os.environ['MINERU_MODEL_SOURCE'] = "modelscope"

    # 4. 默认使用 pipeline 后端
    parse_doc(doc_path_list, output_dir, backend="pipeline")

    # 5. 如需启用 VLM，请取消下面相应注释并填写 server_url
    # parse_doc(doc_path_list, output_dir, backend="vlm-transformers")
    # parse_doc(doc_path_list, output_dir, backend="vlm-sglang-engine")
    # parse_doc(
    #     doc_path_list,
    #     output_dir,
    #     backend="vlm-sglang-client",
    #     server_url="http://127.0.0.1:30000"
    # )